module Language.Coal.AST(transform, Program, Function(..), Block(..), Expression(..), Statement(..)) where

import Language.Coal.Type (Type)
import qualified Language.Coal.Type as Type

import Language.Coal.RawAST (Expression(..))
import qualified Language.Coal.RawAST as RAST

import qualified Language.Coal.VariablesContext as VC

import Language.Coal.Error

import Language.Coal.Util

import Data.Map (Map)
import qualified Data.Map as Map

import Control.Monad.State

import Control.Monad.Error

import Control.Monad

type Program = Map String Function

data Function = Function { args :: [String], body :: Block, ftype :: Type }	
	deriving (Show, Eq, Ord)

data Block = Block {statements :: [Statement], declarations :: Map String Type}
	deriving (Show, Eq, Ord)

data Statement = Expression Expression |
	If Expression Block |
	While Expression Block |
	Return Expression
	deriving (Show, Eq, Ord)

data Context = Context { varInfo :: VC.VariablesContext VariableInformation, freeVariables :: [Type]}
	deriving (Show)

data VariableInformation = VariableInformation {vtype :: Type}
	deriving (Show)

setVarInfo (Context vi fv) nvi = Context nvi fv

insertVariable :: String -> Type -> Context -> Context
insertVariable name t c = setVarInfo c (VC.insertVariable name (VariableInformation t) (varInfo c))

memberVariable :: String -> Context -> Bool
memberVariable name c = VC.memberVariable name (varInfo c)

fromList :: [(String, Type)] -> Context
fromList l = Context (VC.fromList $ mapSnd VariableInformation l) Type.variables

insertVariableList :: [(String, Type)] -> Context -> Context
insertVariableList l c = setVarInfo c (VC.insertVariableList (mapSnd VariableInformation l) (varInfo c))

lookupVariable :: String -> Context -> Maybe VariableInformation 
lookupVariable n c = VC.lookupVariable n (varInfo c) 

child :: Context -> Context
child p = setVarInfo p (VC.child $ varInfo p)

current :: Context -> Map String VariableInformation
current c = VC.current $ varInfo $ c

setFreeVariables :: Context -> [Type] -> Context
setFreeVariables c v = Context vi v where
	vi = varInfo c
	v = freeVariables c

nextVariable :: StateT Context CoalError Type
nextVariable = do
	e <- get
	let r : rest = freeVariables e
	modify $ flip setFreeVariables rest
	return r

addStatement s (Block sl dm) = Block (s : sl) dm

checkVisibility :: String -> StateT Context CoalError ()
checkVisibility name = do 
	e <- get
	unless (memberVariable name e) $ throwError ("Undefenite variable " ++ name)

transformType :: Maybe Type -> StateT Context CoalError Type
transformType Nothing = nextVariable
transformType (Just t) = return t
                                                   
transform :: [(String, Type)] -> [RAST.Function] -> CoalError Program

transform defaultContext program = evalStateT (transformProgram program) contextWithFunctions  where
	functionsDeclarations = map transformFunctionDeclare program
	contextWithFunctions = insertVariableList functionsDeclarations (child $ fromList defaultContext)

transformProgram :: [RAST.Function] -> StateT Context CoalError Program
transformProgram program = do
	tprogram <- mapM transformFunction program
	return $ Map.fromList $ zip (map RAST.name program) tprogram where	

transformFunction :: RAST.Function -> StateT Context CoalError Function

transformFunction f = do
	e <- get
	modify (child . insertVariableList (RAST.args f))
	tbody <- transformBlock $ RAST.body f
	me <- get
	put $ setFreeVariables e (freeVariables me)
	return $ Function argsNames tbody functionType where
	argsNames = map fst (RAST.args f)
	functionType = Type.function (map snd (RAST.args f)) (RAST.rtype f)

transformFunctionDeclare :: RAST.Function -> (String, Type)

transformFunctionDeclare f = (RAST.name f, Type.function (map snd $ RAST.args f) (RAST.rtype f))

transformBlock :: [RAST.Statement] -> StateT Context CoalError Block
transformBlock sl = do
	e <- get
	put $ child e
	tsl <- transformBlock' sl	
	me <- get
	put $ setFreeVariables e (freeVariables me)
	return tsl

transformBlock' :: [RAST.Statement] -> StateT Context CoalError Block

transformBlock' [] = do
	e <- get
	return (Block [] (Map.map vtype (current e))) 

transformBlock' (RAST.Expression expr : rest) = do
	texpr <- transformExpression expr
	trest <- transformBlock' rest
	return $ addStatement (Expression texpr) trest

transformBlock' (RAST.Return expr : rest) = do
	texpr <- transformExpression expr
	trest <- transformBlock' rest
	return $ addStatement (Return texpr) trest

transformBlock' (RAST.DeclareAndAsign n expr t : rest) = do
	texpr <- transformExpression expr
	tt <- transformType t
	modify $ insertVariable n tt
	trest <- transformBlock' rest
	return $ addStatement (Expression $ Asign n texpr) trest

transformBlock' (RAST.Declare n t : rest) = do
	tt <- transformType t
	modify $ insertVariable n tt
	transformBlock' rest

transformBlock' (RAST.If expr current : rest) = do
	texpr <- transformExpression expr
	tcurrent <- transformBlock current
	trest <- transformBlock' rest
	return $ addStatement (If texpr tcurrent) trest

transformBlock' (RAST.While expr current : rest) = do
	texpr <- transformExpression expr
	tcurrent <- transformBlock current
	trest <- transformBlock' rest
	return $ addStatement (While texpr tcurrent) trest


transformExpression :: RAST.Expression -> StateT Context CoalError Expression

transformExpression expr@(Identificator name) = do
	checkVisibility name
	t <- gets $ fromJust . lookupVariable name
	if Type.isFunction $ vtype t then return (Call name []) else return (Identificator name)

transformExpression (Call name args) = do
	checkVisibility name
	targs <- mapM transformExpression args
	return (Call name targs)

transformExpression expr@(Integer i) = return (Integer i)

transformExpression (Asign name expr) = do
	checkVisibility name
	texpr <- transformExpression expr
	return (Asign name texpr)
                                
